#!/usr/bin/env python
# -*- coding: utf-8 -*-

# Copyright 2012 Tobias Marschall
# 
# This file is part of CLEVER.
# 
# CLEVER is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# CLEVER is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with CLEVER.  If not, see <http://www.gnu.org/licenses/>.

from __future__ import print_function, division
from optparse import OptionParser
import sys
import os
import gzip
import subprocess

__author__ = "Tobias Marschall"

usage = """%prog [options] <output-prefix>

Reads alignment priors (e.g. as generated by bam-to-alignment-priors) from 
stdin and writes them to a separate file for each chromosome named
<output-prefix>.<chromosome>.aln-priors."""

def main():
	parser = OptionParser(usage=usage)
	parser.add_option("-z", action="store_true", dest="zip_output", default=False,
					  help="GZIP output files")
	parser.add_option("-c", action="store", dest="chromosomes", default=None,
					  help="Only extracts given chromosomes. Expects comma-separated list.")
	(options, args) = parser.parse_args()
	if (len(args)!=1) or (os.isatty(0)):
		parser.print_help()
		sys.exit(1)
	if options.zip_output:
		zip_executable = None
		if os.system('which gzip > /dev/null 2>&1') == 0:
			zip_executable = 'gzip'
		if os.system('which pigz > /dev/null 2>&1') == 0:
			zip_executable = 'pigz'
	# dictionary mapping chromosome names to output files
	output_files = dict()
	def open_file(chromosome):
		if options.zip_output:
			if zip_executable == None:
				output_files[chromosome] = gzip.open('%s.%s.aln-priors.gz'%(args[0],chromosome), 'wb')
			else:
				output_files[chromosome] = subprocess.Popen(zip_executable, stdin=subprocess.PIPE, stdout=open('%s.%s.aln-priors.gz'%(args[0],chromosome),'w')).stdin
		else:
			output_files[chromosome] = file('%s.%s.aln-priors'%(args[0],chromosome),'w')
	if options.chromosomes != None:
		for chromosome in options.chromosomes.strip().split(','):
			open_file(chromosome)
	for line in sys.stdin:
		fields = line.split()
		assert fields[3] == fields[8]
		chromosome = fields[3]
		if not output_files.has_key(chromosome):
			if options.chromosomes==None: 
				open_file(chromosome)
			else:
				continue
		output_files[chromosome].write(line)

if __name__ == '__main__':
	sys.exit(main())
